#include "gaussian.h"
#define DIM_BLOCK 256

/**
 * Device function so it is not possible to call it from the host. It return the evaluation of a given gaussian in a given point (x, y)
 *
 * @param parameters Gaussian parameters
 * @param x first coordinate of the point
 * @param y second coordinate of the point
 * @return evaluation of gaussian point (x, y)
 */

__device__ float evaluateGaussian (float* parameters, int x, int y) {
    float slope = parameters[PAR_a] * x + parameters[PAR_b] * y + parameters[PAR_c];
    float x_arg = pow (((float) x - parameters[PAR_X]), (float) 2.0) / pow (parameters[PAR_SX], (float) 2.0);
    float y_arg = pow (((float) y - parameters[PAR_Y]), (float) 2.0) / pow (parameters[PAR_SY], (float) 2.0);
    float arg = - (x_arg + y_arg);
    return parameters[PAR_A] * exp (arg) + slope;
}

/**
 * Computes the vector of the image
 *
 * @param image image representing the Gaussian shape
 * @param vetor resultimg gradient vector
 * @param parameters Gaussian paraemters
 * @param width width (and heigth) of the image
 * @param dim dimnesion of the image
 */

__global__ void computeVector (unsigned char* image, float* vector, float* parameters, int width, int dim) {

    int pos, x, y;
    int i = 0, index = threadIdx.x;
    
    /* we allocate Gaussian parameters in shared memory */
    __shared__ float par [DIM_PAR];
    
    /* the first thread will load the parameters */
    if (index == 0)
        for (i = 0; i < DIM_PAR; ++i)
            par[i] = parameters[i];
    
    /* we calculate right index and the coordinates of the point */
    pos = blockIdx.x * DIM_BLOCK + index;
    
    /* if the calculated index is rigth we work */
    if (pos < dim ) {
    
        x = (pos + 1) % width;
        y = (pos + 1) / width;
        vector[pos] = image[pos] - evaluateGaussian (&par[PAR_A], x, y);
    }
    
}

/**
 * Computes the matrix of the image
 *
 * @param image image representing the Gaussian shape
 * @param matrix resultimg gradient matrix
 * @param parameters Gaussian paraemters
 * @param width width (and heigth) of the image
 * @param dim dimnesion of the image
 */

__global__ void computeMatrix (float* matrix, float* parameters, int width, int dim) {
    
    int pos, x , y;
    int i = 0, index = threadIdx.x;
    float dexp, diff_x, diff_y, sig2x, sig2y, amp;
    
    /* we allocate Gaussian parameters in shared memory */
    __shared__ float par [DIM_PAR];
    
    /* the first thread will load the parameters */
    if (index == 0)
        for (i = 0; i < DIM_PAR; ++i)
            par[i] = parameters[i];
    
    /* we calculate right index and the coordinates of the point */
    pos = blockIdx.x * DIM_BLOCK + index;
    
    /* if the calculated index is rigth we work */
    if (pos < dim ) {
        
        x = (pos + 1) % width;
        y = (pos + 1) / width;
        
        diff_x = (float) x - par[PAR_X];
        diff_y = (float) y - par[PAR_Y];
        sig2x = pow (par[PAR_SX], (float) 2.0);
        sig2y = pow (par[PAR_SY], (float) 2.0);
        dexp = exp ( ( (pow (diff_x, (float) 2.0)) / sig2x ) + ( (pow (diff_y, (float) 2.0)) / sig2y ));
        amp = (float) 2.0 * par[PAR_A];
        
        pos = pos * 8;

        matrix[pos] = 1 / dexp;
        matrix[pos + 1] = ( amp * ( (float) x  - par[PAR_X] ) ) / (sig2x * dexp);
        matrix[pos + 2] = ( amp * ( (float) y  - par[PAR_Y] ) ) / (sig2y * dexp);
        matrix[pos + 3] = ( amp * pow (diff_x , (float) 2.0) ) / ( pow (par[PAR_SX], (float) 3.0) * dexp );
        matrix[pos + 4] = ( amp * pow (diff_y , (float) 2.0) ) / ( pow (par[PAR_SY], (float) 3.0) * dexp );
        matrix[pos + 5] = (float) x;
        matrix[pos + 6] = (float) y;
        matrix[pos + 7] = (float) 1.0;
        
    }
    
}

/**
 * Updates the Gaussian parameters
 *
 * @param delta vector with updates
 * @param parameters Gaussian parameters
 */

__global__ void updateParameters (float* delta, float* parameters) {
    int i;
    for (i = 0; i < 8; ++i)
        parameters[i] = parameters[i] + delta[i]; 
}

/**
 * External C code
 */

extern "C" {

    /**
     * Calls the kernel for computing the gradient vector of the image
     *
     * @param image image representing the Gaussian shape
     * @param vetor resultimg gradient vector
     * @param parameters Gaussian paraemters
     * @param width width (and heigth) of the image
     */
    
    void computeVectorWrapper (unsigned char* image, float* vector, float* parameters, int width, int dim) {
        /* we have width * width elements so we calculate the MAXIMUM number of blocks in this way */
        int num_blocks = ( dim + DIM_BLOCK - 1 ) / DIM_BLOCK;
        /* we call a kernel with a grid of num_blocks blocks every one of dimension DIM_BLOCK_SUBVECTOR */
        computeVector<<<num_blocks, DIM_BLOCK>>> (image, vector, parameters, width, dim);
    }
    
    /**
     * Calls the kernel for computing the matrix vector of the image
     *
     * @param matrix resultimg gradient matrix
     * @param parameters Gaussian paraemters
     * @param width width (and heigth) of the image
     */
    
    void computeMatrixWrapper (float* matrix, float* parameters, int width, int dim) {
        /* we have width * width elements so we calculate the MAXIMUM number of blocks in this way */
        int num_blocks = ( dim + DIM_BLOCK - 1 ) / DIM_BLOCK;
        /* we call a kernel with a grid of num_blocks blocks every one of dimension DIM_BLOCK_SUBVECTOR */
        computeMatrix<<<num_blocks, DIM_BLOCK>>> (matrix, parameters, width, dim);
    }
    
    /**
     * Calls the kernel for updating the Gaussian parameters
     *
     * @param delta vector with updates
     * @param parameters Gaussian parameters
     */
    
    void updateParametersWrapper (float* delta, float* parameters) {
        updateParameters <<<1, 1>>> (delta, parameters);
    }
    
}